import argparse
import base64
import io
import json
from pathlib import Path
import sys

import cv2
import numpy as np
import uvicorn
from fastapi import FastAPI, Form, UploadFile, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from rapidocr_onnxruntime import RapidOCR

from base.pdf_to_images import pdf_to_images_import

sys.path.append(str(Path(__file__).resolve().parent.parent))


class OCRAPIUtils:
    def __init__(self) -> None:
        self.ocr = RapidOCR()

    def __call__(self, img):
        img = np.array(img)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        ocr_res, _ = self.ocr(img)
        if not ocr_res:
            return json.dumps({})

        out_dict = {
            str(i): {"rec_txt": rec, "dt_boxes": dt_box, "score": score}
            for i, (dt_box, rec, score) in enumerate(ocr_res)
        }
        return out_dict


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源访问，可以根据需要指定特定的域名
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有 HTTP 方法
    allow_headers=["*"],  # 允许所有请求头
)
processor = OCRAPIUtils()

# ocr识别图片内容
@app.post("/ocr")
async def ocr(image_file: UploadFile = None, image_data: str = Form(None)):
    if image_file:
        img = Image.open(image_file.file)
    elif image_data:
        img_bytes = str.encode(image_data)
        img_b64decode = base64.b64decode(img_bytes)
        img = Image.open(io.BytesIO(img_b64decode))
    else:
        raise ValueError(
            "When sending a post request, data or files must have a value."
        )

    ocr_res = processor(img)
    return ocr_res

# pdf转image
@app.post("/ocr/pdf_to_image")
async def ocr_pdf_to_image(request: Request):
    data = await request.json()

    pdf_url = data.get("pdf_url")

    if pdf_url:
        image_urls = pdf_to_images_import(pdf_url)
        return {"image_urls": image_urls}
    else:
        raise HTTPException(status_code=400, detail="pdf_url is required.")


def main():
    parser = argparse.ArgumentParser("ocr_api")
    parser.add_argument("-ip", "--ip", type=str, default="0.0.0.0", help="IP Address")
    parser.add_argument("-p", "--port", type=int, default=9003, help="IP port")
    args = parser.parse_args()

    cur_file_path = Path(__file__).resolve()
    app_path = f"{cur_file_path.parent.name}.{cur_file_path.stem}:app"
    print(app_path)
    uvicorn.run(app_path, host=args.ip, port=args.port, reload=True)


if __name__ == "__main__":
    main()
